Retain newlines in chat template when `continue_final_message=True` by lewtun · Pull Request #34253 · huggingface/transformers

lewtun · 2024-10-18T21:09:24Z

What does this PR do?

This PR addresses an issue with chat templates producing incorrect inputs when continue_final_message=True and we wish to terminate generation on newlines like \n (as is common in process supervision with math reward models like this one).

In this case, the prefills had their newlines stripped at each step of generation and produced results that disagree with pure greedy decoding to the EOS token. Here's an example:

import copy
from transformers import AutoTokenizer, AutoModelForCausalLM

model_id = "Qwen/Qwen2.5-Math-1.5B-Instruct"
tokenizer = AutoTokenizer.from_pretrained(model_id)
model = AutoModelForCausalLM.from_pretrained(model_id).to("cuda")

### Inputs
system_prompt = "Please reason step by step, and put your final answer within \\boxed{}."
query = "Which number is larger, 9.8 or 9.11?"
messages = [{"role": "system", "content": system_prompt}, {"role": "user", "content": query}]

### Reference - greedy decoding to EOS
inputs = tokenizer.apply_chat_template(messages, return_tensors="pt", return_dict=True, add_generation_prompt=True).to("cuda")
output = model.generate(**inputs, max_new_tokens=256, do_sample=False)
ref_output = tokenizer.decode(output[0][inputs.input_ids.shape[-1]:])
print(f"\n=== Ref completion ===\n{ref_output}")

### Iterative decoding with prefill
current_messages = copy.deepcopy(messages)

for i in range(5):
    if current_messages[-1]["role"] == "user":
        inputs = tokenizer.apply_chat_template(current_messages, return_tensors="pt", return_dict=True, add_generation_prompt=True).to("cuda")
    else:
        inputs = tokenizer.apply_chat_template(current_messages, return_tensors="pt", return_dict=True, continue_final_message=True).to("cuda")
    output = model.generate(**inputs, max_new_tokens=256, do_sample=False, stop_strings=["\n"], tokenizer=tokenizer)
    output_decoded = tokenizer.decode(output[0][inputs.input_ids.shape[-1]:])
    print(f"\n=== Completion {i+1} ===\n{repr(output_decoded)}")
    if current_messages[-1]["role"] == "user":
        current_messages.append({"role": "assistant", "content": output_decoded})
    else:
        current_messages[-1]["content"] += output_decoded

# we strip the final output to remove any trailing newlines produced by terminating on `\n`
final_output = current_messages[-1]["content"].rstrip()
print(f"\n=== Final output ===\n{final_output}")
assert ref_output == final_output

Here's the reference completion from greedy decoding until EOS (note the new lines):

To determine which number is larger between 9.8 and 9.11, we can compare them digit by digit starting from the leftmost digit.

1. Both numbers have the same digit in the units place, which is 9.
2. Next, we compare the digits in the tenths place. The digit in the tenths place of 9.8 is 8, and the digit in the tenths place of 9.11 is 1. Since 8 is greater than 1, we can conclude that 9.8 is greater than 9.11.

Therefore, the larger number is \boxed{9.8}.<|im_end|>

Without the fix we get bad generations because the prefill is continually stripping the trailing new lines per iteration:

=== Completion 1 ===
'To determine which number is larger between 9.8 and 9.11, we can compare them digit by digit starting from the leftmost digit.\n\n'

=== Completion 2 ===
" Here's the step-by-step process:\n\n"

=== Completion 3 ===
' \n\n'

=== Completion 4 ===
' \n\n'

=== Completion 5 ===
' \n\n'

=== Final output ===
To determine which number is larger between 9.8 and 9.11, we can compare them digit by digit starting from the leftmost digit.

 Here's the step-by-step process:

With the proposed fix, order is restored and the resulting completion is sane:

=== Completion 1 ===
'To determine which number is larger between 9.8 and 9.11, we can compare them digit by digit starting from the leftmost digit.\n\n'

=== Completion 2 ===
'1. Both numbers have the same digit in the units place, which is 9.\n'

=== Completion 3 ===
'2. Next, we compare the digits in the tenths place. The digit in the tenths place of 9.8 is 8, and the digit in the tenths place of 9.11 is 1. Since 8 is greater than 1, we can conclude that 9.8 is greater than 9.11.\n\n'

=== Completion 4 ===
'Therefore, the larger number is \\boxed{9.8}.<|im_end|>'

=== Completion 5 ===
'\n'

=== Final output ===
To determine which number is larger between 9.8 and 9.11, we can compare them digit by digit starting from the leftmost digit.

1. Both numbers have the same digit in the units place, which is 9.
2. Next, we compare the digits in the tenths place. The digit in the tenths place of 9.8 is 8, and the digit in the tenths place of 9.11 is 1. Since 8 is greater than 1, we can conclude that 9.8 is greater than 9.11.

Therefore, the larger number is \boxed{9.8}.<|im_end|>

Questions

I am not sure why stripping was included in the original implementation. If it's important to retain, could we instead expose a strip_prefill arg in apply_chat_template that is True by default but can be disabled in cases like process supervision?
I wasn't sure how to write a regression test for this behaviour since the tiny-gpt2 models aren't prone to emitting newlines and running a 1.5B Qwen model in the CI is presumably discouraged.

Before submitting

This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
Did you read the contributor guideline,
Pull Request section?
Was this discussed/approved via a Github issue or the forum? Please add a link
to it if that's the case.
Did you make sure to update the documentation with your changes? Here are the
documentation guidelines, and
here are tips on formatting docstrings.
Did you write any new necessary tests?

Who can review?

Anyone in the community is free to review the PR once the tests have passed. Feel free to tag
members/contributors who may be interested in your PR.

lewtun · 2024-10-18T21:16:43Z

CI seems to be failing for reasons unrelated to this PR

Rocketknight1 · 2024-10-21T12:56:53Z

Wow, this issue was only introduced extremely recently, in #34214 ! I completely didn't foresee this being an issue when I merged that one.

The fundamental problem is that some templates strip() input text, and others don't. This means that the string-matching we do in continue_final_message isn't guaranteed to find a match if we don't also strip the search string. However, as you've noted, this can cause other problems.

Maybe the right solution is to first try to match the input with trailing whitespace, and then fallback to searching for the stripped version if that fails?

lewtun · 2024-10-21T13:00:22Z

Maybe the right solution is to first try to match the input with trailing whitespace, and then fallback to searching for the stripped version if that fails?

I like this approach - let me amend my PR. Do we also need to retain the rstrip() that's applied to rendered_chat?

Rocketknight1 · 2024-10-21T13:06:02Z

No, I don't think so! I think it would break your task if we did, and I really just included it because I didn't realize the trailing whitespace could be meaningful.

Rocketknight1 · 2024-10-21T13:21:28Z

cc @schoennenbeck, who wrote #34214 - you don't need to do anything here, just making you aware that we might change the behaviour again!

lewtun · 2024-10-25T08:51:23Z

@Rocketknight1 I've added a try/except in 0fd98e1

I've also tested that the Llama code still works

# From: https://github.com/huggingface/transformers/pull/34214
from transformers import AutoTokenizer
tokenizer = AutoTokenizer.from_pretrained("hugging-quants/Meta-Llama-3.1-8B-Instruct-AWQ-INT4")
messages=[
    {"role": "user", "content": "What is the capital of France?"},
    {"role": "assistant", "content": "Great question! The capital of France is "}
]
s = tokenizer.apply_chat_template(
    messages, add_generation_prompt=False, 
    continue_final_message=True,
    tokenize=False
)
print(s)

Would you like me to include this as a regression test? I'm also not sure if we need something for the Qwen model, but for that we'd want to run it on GPU

HuggingFaceDocBuilderDev · 2024-10-25T09:13:50Z

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

Rocketknight1 · 2024-10-25T11:55:30Z

@lewtun yeah, this feels like an area where other changes could easily break things again. A regression test would definitely be appreciated!

lewtun · 2024-10-25T14:29:39Z

                )

+    @require_jinja
+    def test_continue_final_message_with_trim(self):


Added a regression test here @Rocketknight1

Rocketknight1

This looks like a great PR to me, and the regression test is nice! I left two nits, but both are just code style rather than anything that actually affects the output.

Rocketknight1 · 2024-10-25T17:21:33Z

+                # Assert that the final message is unterminated
+                self.assertEqual(
+                    prefill_output,
+                    "<|im_start|>system\nsystem message<|im_end|>\n<|im_start|>user\nuser message<|im_end|>\n<|im_start|>assistant\nassistant message",


Nice! This is a clean test.

Co-authored-by: Matt <Rocketknight1@users.noreply.github.com>

lewtun · 2024-10-27T06:08:22Z

@Rocketknight1 I've included your suggestions and the CI seems to be failing for reasons unrelated to this PR 😢 . What do you recommend doing in this case?

Rocketknight1 · 2024-10-29T18:02:59Z

@lewtun CI looks green now, so I'm happy! cc @ArthurZucker / @LysandreJik for core maintainer review and we should be good to go.

LysandreJik

Thanks @lewtun for the PR and @Rocketknight1 for the review!

…uggingface#34253) * Retain newlines in chat template when * Add try/except * Add regression test * Simplify test * Apply suggestions from code review Co-authored-by: Matt <Rocketknight1@users.noreply.github.com> --------- Co-authored-by: Matt <Rocketknight1@users.noreply.github.com>

Retain newlines in chat template when

3dcaee1

lewtun requested a review from Rocketknight1 October 18, 2024 21:09

lewtun added 2 commits October 25, 2024 10:44

Merge branch 'main' into fix-continue-msg

aeeee98

Add try/except

0fd98e1

lewtun commented Oct 25, 2024

View reviewed changes

Comment thread src/transformers/tokenization_utils_base.py Resolved

Merge branch 'main' into fix-continue-msg

92fed23

lewtun added 3 commits October 25, 2024 15:31

Merge branch 'main' into fix-continue-msg

39d4f42

Add regression test

138ee7b

Simplify test

2cbbc1f

lewtun commented Oct 25, 2024

View reviewed changes

Merge branch 'main' into fix-continue-msg

24ac783

Rocketknight1 approved these changes Oct 25, 2024

View reviewed changes

lewtun commented Oct 26, 2024

View reviewed changes

Comment thread src/transformers/tokenization_utils_base.py Resolved

lewtun and others added 2 commits October 26, 2024 13:20

Apply suggestions from code review

a1f99aa

Co-authored-by: Matt <Rocketknight1@users.noreply.github.com>

Merge branch 'main' into fix-continue-msg

242a073

Merge branch 'main' into fix-continue-msg

4d00a03

lewtun added 2 commits November 4, 2024 17:30

Merge branch 'main' into fix-continue-msg

a245df1

Merge branch 'main' into fix-continue-msg

0f05f9e

LysandreJik approved these changes Nov 15, 2024

View reviewed changes

Rocketknight1 merged commit 8ba3e15 into main Nov 15, 2024

Rocketknight1 deleted the fix-continue-msg branch November 15, 2024 14:27

Conversation

lewtun commented Oct 18, 2024 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

What does this PR do?

Before submitting

Who can review?

Uh oh!

lewtun commented Oct 18, 2024

Uh oh!

Rocketknight1 commented Oct 21, 2024 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

lewtun commented Oct 21, 2024

Uh oh!

Rocketknight1 commented Oct 21, 2024

Uh oh!

Rocketknight1 commented Oct 21, 2024

Uh oh!

Uh oh!

lewtun commented Oct 25, 2024

Uh oh!

HuggingFaceDocBuilderDev commented Oct 25, 2024

Uh oh!

Rocketknight1 commented Oct 25, 2024

Uh oh!

lewtun Oct 25, 2024

Choose a reason for hiding this comment

Uh oh!

Rocketknight1 left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Rocketknight1 Oct 25, 2024

Choose a reason for hiding this comment

Uh oh!

Uh oh!

lewtun commented Oct 27, 2024

Uh oh!

Rocketknight1 commented Oct 29, 2024

Uh oh!

LysandreJik left a comment

Choose a reason for hiding this comment

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

lewtun commented Oct 18, 2024 •

edited

Loading

Rocketknight1 commented Oct 21, 2024 •

edited

Loading